import os
import random
import shutil
from tqdm import tqdm
import numpy as np

def remove_dir(path):
    try:
        shutil.rmtree(path)
    except:
        pass


def label_of_directory(directory):
    """
    sorted for label indices
    return a dict for {'classes', 'range(len(classes))'}
    """
    classes = []
    for subdir in sorted(os.listdir(directory)):
        if os.path.isdir(os.path.join(directory, subdir)):
            classes.append(subdir)
    class_indices = dict(zip(classes, range(len(classes))))
    return class_indices


def write_txt(path, type):
    dict = label_of_directory(os.path.join(path, "{}_data".format(type)))
    alllist = []
    index = 0
    for folder in dict:
        img_list = [f for f in os.listdir(os.path.join(path, "{}_data".format(type), folder)) if not f.startswith('.')]
        for img in img_list:
            str0 = '%d\t%s\t%d\n' % (index, folder + '/' + img, dict[folder])
            index += 1
            alllist.append(str0)
    file = open(os.path.join(path, "{}.txt".format(type)), "w")
    for str1 in alllist:
        file.write(str1)
    file.close()


def convert_images(converted_path, path, split_ratio):
    train_path = os.path.join(converted_path, 'train_data')
    test_path = os.path.join(converted_path, 'test_data')
    valid_path = os.path.join(converted_path, 'valid_data')

    os.mkdir(train_path)
    os.mkdir(test_path)
    os.mkdir(valid_path)

    for root, dirs, files in os.walk(path):
        if root == path:
            continue
        category = os.path.basename(root)
        label = category
        jpgpath_train = (os.path.join(train_path, str(label)))
        jpgpath_valid = (os.path.join(valid_path, str(label)))
        jpgpath_test = (os.path.join(test_path, str(label)))
        os.mkdir(jpgpath_train)
        os.mkdir(jpgpath_valid)
        os.mkdir(jpgpath_test)

        random.shuffle(files)  # 打乱数据集
        count = 0
        for name in tqdm(files):
            if count < int(len(files) * split_ratio[0]):  # 放到训练集
                shutil.copy(os.path.join(root, name), os.path.join(jpgpath_train, name))
            elif count < int(len(files) * (split_ratio[0] + split_ratio[1])):  # 放到测试集
                shutil.copy(os.path.join(root, name), os.path.join(jpgpath_valid, name))
            else:
                shutil.copy(os.path.join(root, name), os.path.join(jpgpath_test, name))
            count += 1

    write_txt(converted_path, 'train')
    write_txt(converted_path, 'valid')
    write_txt(converted_path, 'test')


if __name__ == '__main__':

    imagesPath = r'E:\数据集\RGB_15_train'  # 数据集路径
    """
    class1
    class2
    class3
    ...
    """
    converted_path = 'GID'  # 数据保存路径

    train_split_ratio = 0.8  # ratio of train set size
    valid_split_ratio = 0.2  # ratio of valid set size
    test_split_ratio = 0  # ratio of test set size

    split_ratio = [train_split_ratio, valid_split_ratio, test_split_ratio]

    if np.sum(np.array(split_ratio)) != 1:
        raise AssertionError("The sum of the ratio of train set, valid set and test set is not equal to 1")
    if split_ratio[0] == 0:
        raise AssertionError("The ratio of train set size should be greater than 0")
    if os.path.exists(converted_path):
        remove_dir(converted_path)
    os.mkdir(converted_path)
    convert_images(converted_path, imagesPath, split_ratio)
